import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from std_msgs.msg import Int32MultiArray 
import cv2
import numpy as np

class CvCone(Node):
    def __init__(self):
        super().__init__('cv_cone')
        self.subscription = self.create_subscription(
            Image,
            '/image',
            self.listener_callback,
            10)
        self.cone_publisher = self.create_publisher(Int32MultiArray, '/cone', 10)  # 创建新的发布者，用于发布x和y坐标
        self.up = 120
        self.left = 50
        self.area = 1200

    def listener_callback(self, msg):
        np_arr = np.frombuffer(msg.data, np.uint8)
        cv_image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
        x, y, w= self.process_frame(cv_image)

        cone_msg = Int32MultiArray()
        cone_msg.data = [x, y, w]

        self.cone_publisher.publish(cone_msg)

    
    def process_frame(self,frame):
        x0 = 0
        y0 = 0
        w0 = 0
        im0 = frame[self.up:450, self.left:600]
        hsv_img = cv2.cvtColor(im0, cv2.COLOR_BGR2HSV)

        # Define range of target color in HSV
        lower_hsv1 = np.array([0, 0, 0])
        lower_hsv2 = np.array([82, 255, 255])
        upper_hsv1 = np.array([0, 0, 0])
        upper_hsv2 = np.array([255, 90, 255])

        # Threshold the HSV image
        imgThreshLow = cv2.inRange(hsv_img, lower_hsv1, lower_hsv2)
        imgThreshHigh = cv2.inRange(hsv_img, upper_hsv1, upper_hsv2)
        threshed_img = cv2.bitwise_or(imgThreshLow, imgThreshHigh)
        '''
        # Noise removal
        kernel = np.ones((5, 5), np.uint8)
        smoothed_img = cv2.dilate(threshed_img, kernel, iterations=11)
        smoothed_img = cv2.erode(smoothed_img, kernel, iterations=7)
        '''
        # Find contours in the black regions of smoothed_img
        contours, hierarchy = cv2.findContours(255 - threshed_img, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

        # Parameters for text display
        font = cv2.FONT_HERSHEY_SIMPLEX
        fontScale = 0.6
        fontColor = (0, 0, 255)
        lineType = 2

        # Initialize variables to store the largest contour
        max_area = 0
        largest_rect = None
        largest_cnt = None

        # Identify the largest contour and store its bounding rectangle
        for cnt in contours:
            area = cv2.contourArea(cnt)
            if area > self.area:  # Filter small contours if needed
                x, y, w, h = cv2.boundingRect(cnt)
                if area > max_area:
                    max_area = area
                    largest_rect = (x, y, w, h)
                    largest_cnt = cnt

        # Draw bounding rectangles and centroids
        for cnt in contours:
            area = cv2.contourArea(cnt)
            if area > self.area:
                x, y, w, h = cv2.boundingRect(cnt)

                # Draw the largest rectangle in red, others in green
                if (x, y, w, h) == largest_rect:
                    color = (0, 0, 255)  # Red for the largest rectangle
                    x0 = int(x + self.left + w / 2)
                    y0 = int(y + self.up + h )
                    w0 = int(w)
                else:
                    color = (0, 255, 0)  # Green for other rectangles
                #cv2.rectangle(frame, (x + self.left, y + self.up), (x + w + self.left, y + h + self.up), color, 3)
                '''
                # Calculate the center (centroid) of the contour
                M = cv2.moments(cnt)
                if M["m00"] != 0:
                    cx = int(M["m10"] / M["m00"]) + 50
                    cy = int(M["m01"] / M["m00"]) + 120
                    cv2.circle(frame, (cx, cy), 5, (255, 0, 0), -1)
                    cv2.putText(frame, 'Cone Center', (cx, cy - 10), font, fontScale, fontColor, lineType)
                '''
        return x0,y0,w0

def main(args=None):
    rclpy.init(args=args)
    cv_cone = CvCone()
    rclpy.spin(cv_cone)
    rclpy.shutdown()

if __name__ == '__main__':
    main()